import argparse
import logging
import re
import os
import sys
import json

import numpy as np
import scipy.io
import pyulog


def _sanitize_varname(name: str) -> str:
    # 将话题名转换为 MATLAB 合法变量名风格（字母数字和下划线）
    s = re.sub(r"[^0-9a-zA-Z_]", "_", name)
    if not s:
        s = "topic"
    if s[0].isdigit():
        s = "t_" + s
    return s


def _sanitize_fieldname(name: str, existing: set, max_len: int = 31) -> str:
    # 清理字段名（只保留字母数字和下划线），并截断以满足 MATLAB 的 31 字符限制，
    # 若发生冲突在尾部添加序号（仍保证不超过限制）。
    base = re.sub(r"[^0-9a-zA-Z_]", "_", name)
    if not base:
        base = "f"
    # 保留空间用于后缀（如 _1），初始尝试不带后缀
    candidate = base[:max_len]
    if candidate not in existing:
        existing.add(candidate)
        return candidate

    # 如果冲突，添加数字后缀并确保总长度不超过 max_len
    i = 1
    while True:
        suffix = f"_{i}"
        avail = max_len - len(suffix)
        cand = (base[:avail] + suffix)
        if cand not in existing:
            existing.add(cand)
            return cand
        i += 1


def _prepare_for_savemat(src: dict) -> dict:
    """
    将 mat 字典转换为 savemat 可以接受的形式：
    - numpy structured arrays 转换为 dict（field -> array），并保证字段名 <= 31 且唯一
    - None 替换为空数组
    - 嵌套结构也确保字段名符合要求
    """
    def _convert_struct_array(arr, parent_existing=None):
        # 处理结构化数组，确保所有层级的字段名都符合要求
        if parent_existing is None:
            parent_existing = set()
            
        if arr is None:
            return np.array([], dtype=object)
            
        # 处理 numpy 结构化数组
        if isinstance(arr, (np.ndarray, np.void)) and getattr(getattr(arr, 'dtype', None), 'names', None):
            existing = parent_existing.copy()
            struct_dict = {}
            
            for fname in arr.dtype.names:
                safe = _sanitize_fieldname(fname, existing, max_len=31)
                try:
                    if isinstance(arr, np.void):
                        # 单条记录
                        field_data = np.array([arr[fname]])
                    else:
                        # 数组
                        field_data = arr[fname]
                        
                    # 如果字段值本身是结构体，递归处理
                    if isinstance(field_data, (np.ndarray, np.void)) and \
                       getattr(getattr(field_data, 'dtype', None), 'names', None):
                        field_data = _convert_struct_array(field_data, existing)
                    elif np.isscalar(field_data):
                        field_data = np.array([field_data])
                        
                    struct_dict[safe] = field_data
                except Exception as e:
                    logging.warning(f"处理字段 {fname} 时出错: {e}")
                    struct_dict[safe] = np.array([], dtype=object)
            
            return struct_dict
            
        # 处理字典
        if isinstance(arr, dict):
            existing = parent_existing.copy()
            result = {}
            for k, v in arr.items():
                safe = _sanitize_fieldname(k, existing, max_len=31)
                result[safe] = _convert_struct_array(v, existing)
            return result
            
        # 处理列表
        if isinstance(arr, (list, tuple)):
            try:
                return np.array(arr)
            except Exception:
                return np.array(arr, dtype=object)
                
        return arr

    # 主处理逻辑
    out = {}
    top_existing = set()
    
    for k, v in src.items():
        # 顶层键名也需要处理
        safe_key = _sanitize_fieldname(k, top_existing, max_len=31)
        out[safe_key] = _convert_struct_array(v)
        
    return out


def _extract_payload(obj):
    # 尝试从 pyulog 返回的对象中提取可序列化的数据
    # 常见属性包括 `.data`（numpy structured array 或 list）、`.field_data`、`.get_field_data()`
    if obj is None:
        return None
    if hasattr(obj, "data"):
        payload = obj.data
    elif hasattr(obj, "field_data"):
        payload = obj.field_data
    elif hasattr(obj, "get_field_data"):
        try:
            payload = obj.get_field_data()
        except Exception:
            payload = obj
    else:
        payload = obj

    # 将 Python 列表尽量转换为 numpy 数组，失败时保存为 object 数组
    if isinstance(payload, list):
        try:
            payload = np.array(payload)
        except Exception:
            payload = np.array(payload, dtype=object)
    return payload


def load_topic_mapping(json_path: str) -> dict:
    """
    从 JSON 文件加载话题映射配置
    JSON 格式示例:
    {
        "vehicle_local_position": "local_pos",
        "sensor_combined": "imu"
    }
    """
    if not os.path.isfile(json_path):
        raise FileNotFoundError(f"找不到话题映射文件: {json_path}")
    
    with open(json_path, 'r') as f:
        mapping = json.load(f)
    
    if not isinstance(mapping, dict):
        raise ValueError("JSON 文件必须包含话题映射字典")
    
    return mapping

def extract_topics_from_ulog(ulog_file_path: str, output_mat_file_path: str, 
                           topic_mapping: dict = None, verbose: bool = False):
    if verbose:
        logging.basicConfig(level=logging.INFO)
    else:
        logging.basicConfig(level=logging.WARNING)

    if not os.path.isfile(ulog_file_path):
        raise FileNotFoundError(f"ULog 文件不存在: {ulog_file_path}")

    # 读取 ULog
    ulog = pyulog.ULog(ulog_file_path)

    # 取得可用的话题列表
    if hasattr(ulog, "get_topic_names"):
        available_topics = ulog.get_topic_names()
    else:
        available_topics = [m.name for m in getattr(ulog, "data_list", [])]

    # 确定要处理的话题
    if topic_mapping:
        # 只处理映射中指定的话题
        topics_to_process = [t for t in topic_mapping.keys() if t in available_topics]
        if verbose:
            logging.info(f"找到 {len(available_topics)} 个话题，将处理 {len(topics_to_process)} 个映射话题")
    else:
        # 处理所有话题
        topics_to_process = available_topics
        if verbose:
            logging.info(f"找到 {len(topics_to_process)} 个话题")

    mat_dict = {}
    for topic in topics_to_process:
        try:
            # 尝试通过 API 获取话题数据对象
            if hasattr(ulog, "get_topic_data"):
                obj = ulog.get_topic_data(topic)
            else:
                obj = next((m for m in getattr(ulog, "data_list", []) if m.name == topic), None)

            payload = _extract_payload(obj)

            # 使用映射名称或原始话题名
            if topic_mapping and topic in topic_mapping:
                key = _sanitize_varname(topic_mapping[topic])
            else:
                key = _sanitize_varname(topic)

            # 避免同名覆盖
            base = key
            i = 1
            while key in mat_dict:
                key = f"{base}_{i}"
                i += 1

            mat_dict[key] = payload
            if verbose:
                logging.info(f"添加话题: {topic} -> {key} (type={type(payload)})")
        except Exception as e:
            logging.warning(f"跳过话题 {topic}: {e}")

    # 保存为 .mat
    savable = _prepare_for_savemat(mat_dict)
    scipy.io.savemat(output_mat_file_path, savable, do_compression=False)
    print(f"数据已保存到 {output_mat_file_path}")

def main():
    parser = argparse.ArgumentParser(description="从 ULog (.ulg) 文件提取 uORB 话题并保存为 .mat 文件")
    parser.add_argument("ulog_file", type=str, help="输入的 ULog (.ulg) 文件路径")
    parser.add_argument("output_mat_file", type=str, help="输出的 .mat 文件路径")
    parser.add_argument("-v", "--verbose", action="store_true", help="输出详细信息")
    parser.add_argument("--list", action="store_true", help="仅列出可用的话题，不保存")
    parser.add_argument("--topic-map", type=str, help="话题映射 JSON 文件路径")

    args = parser.parse_args()

    if not os.path.isfile(args.ulog_file):
        print(f"错误：找不到文件 {args.ulog_file}", file=sys.stderr)
        sys.exit(2)

    # 如果只列出话题
    if args.list:
        ulog = pyulog.ULog(args.ulog_file)
        if hasattr(ulog, "get_topic_names"):
            topics = ulog.get_topic_names()
        else:
            topics = [m.name for m in getattr(ulog, "data_list", [])]
        for t in topics:
            print(t)
        return

    # 加载话题映射（如果提供）
    topic_mapping = None
    if args.topic_map:
        try:
            topic_mapping = load_topic_mapping(args.topic_map)
        except Exception as e:
            print(f"加载话题映射失败: {e}", file=sys.stderr)
            sys.exit(3)

    extract_topics_from_ulog(args.ulog_file, args.output_mat_file, 
                           topic_mapping=topic_mapping, verbose=args.verbose)

if __name__ == "__main__":
    main()
